package com.zws.cucumber.adaptation.cass;

import com.datastax.oss.driver.api.core.*;
import com.datastax.oss.driver.api.core.context.DriverContext;
import com.datastax.oss.driver.api.core.cql.*;
import com.datastax.oss.driver.api.core.metadata.Metadata;
import com.datastax.oss.driver.api.core.metrics.Metrics;
import com.datastax.oss.driver.api.core.session.Request;
import com.datastax.oss.driver.api.core.type.reflect.GenericType;
import com.datastax.oss.driver.internal.core.cql.Conversions;
import com.datastax.oss.driver.internal.core.cql.ResultSets;
import com.zws.cucumber.util.cass.ZwsCassTemplate;
import com.zws.cucumber.util.cass.ZwsCassTemplateImpl;
import com.zws.cucumber.util.colorfullog.ZwsColorfulLogger;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import org.slf4j.Logger;

import java.util.*;
import java.util.concurrent.*;


/**
 * author: zws
 */
public class ZwsCqlSessionProxy implements CqlSession {

    private static final Logger logger = ZwsColorfulLogger.of(ZwsCqlSessionProxy.class);

    public static final DriverException EX_TIMEOUT = new DriverTimeoutException("Query timed out after PT2S");
    public static final DriverException EX_NO_NODE_AVAILABLE = new NoNodeAvailableException();

    private static ZwsCqlSessionProxy instance;
    private static volatile boolean needThrowException = false;
    private static volatile DriverException targetExceptionToThrow = null;
    private static volatile String targetTableToThrowException = null;
    private static volatile ZwsCqlDmlType targetDmlTypeToThrowException = null;
    private static volatile Future<?> asyncExecutionFuture = null;
    private static volatile BlockingQueue<Future<?>> asyncExecutionFutureQueue = new ArrayBlockingQueue<>(1024);
    /* don't use */
    private static boolean batching = false;
    /* don't use */
    private static BatchStatement batchStmt = null;
    private static ExecutorService executorService = Executors.newFixedThreadPool(1);
    private static volatile int maxExecutionId = 0;

    public static void enableThrowingException(DriverException exception, String table, ZwsCqlDmlType dmlType) {
        ZwsCqlSessionProxy.needThrowException = true;
        targetExceptionToThrow = exception;
        targetTableToThrowException = table;
        targetDmlTypeToThrowException = dmlType;
    }

    public static void enableThrowingException(DriverException exception, String table) {
        enableThrowingException(exception, table, null);
    }

    public static void enableThrowingException(DriverException exception, ZwsCqlDmlType dmlType) {
        enableThrowingException(exception, null, dmlType);
    }

    public static void enableThrowingException(DriverException exception) {
        enableThrowingException(exception, null, null);
    }

    public static void enableThrowingException() {
        enableThrowingException(null, null, null);
    }

    public static void disableThrowingException() {
        needThrowException = false;
        targetExceptionToThrow = null;
        targetTableToThrowException = null;
        targetDmlTypeToThrowException = null;
    }

    private static int nextExecutionId() {
        return ++maxExecutionId;
    }

    public static void asyncTruncate(ZwsCassTemplate cassTemplate, String keyspace) {
        waitAsyncExecution();
        int id = nextExecutionId();
        asyncExecutionFutureQueue.offer(executorService.submit(() -> {
            long s = System.currentTimeMillis();
            cassTemplate.cleanData(true, keyspace);
            logger.debug("async truncate task[{}] finished, elapsed: {}ms", id, System.currentTimeMillis() - s);
        }));
        logger.debug("submit async truncate task[{}]", id);
    }

    public static void asyncExecute(Iterable<String> modifiedTables, List<String> qlList) {
        int id = nextExecutionId();
        asyncExecutionFutureQueue.offer(executorService.submit(() -> {
            logger.debug("Execute Batch Statement {}", qlList);
            ZwsCassTemplateImpl.addModifiedTables(modifiedTables);

            long s = System.currentTimeMillis();
            BatchStatement batchStatement = BatchStatement.newInstance(BatchType.LOGGED);
            for (String cql : qlList) {
                SimpleStatement simpleStatement = SimpleStatement.newInstance(cql);
                if (batchStatement.computeSizeInBytes(instance.getPuppet().getContext()) + simpleStatement.computeSizeInBytes(instance.getPuppet().getContext()) + 512 > 51200) {
                    instance.getPuppet().execute(batchStatement);
                    batchStatement = BatchStatement.newInstance(BatchType.LOGGED);
                }
                batchStatement = batchStatement.add(simpleStatement);
            }
            if (batchStatement.size() > 0) {
                instance.getPuppet().execute(batchStatement);
            }
            long e = System.currentTimeMillis();

            logger.debug("async execution task[{}] finished, elapsed: {}ms", id, e - s);
        }));
        logger.debug("submit async execution task[{}]", id);
    }

    public static void asyncExecute0(Collection<String> modifiedTables, List<? extends BatchableStatement> batchableStatementList) {
        int id = nextExecutionId();
        asyncExecutionFutureQueue.offer(executorService.submit(() -> {
            long s = System.currentTimeMillis();
            logger.debug("async execution task[{}] start", id);
            ZwsCassTemplateImpl.addModifiedTables(modifiedTables);

            BatchStatement batchStatement = BatchStatement.newInstance(BatchType.LOGGED);
            for (BatchableStatement stmt : batchableStatementList) {
                if (batchStatement.computeSizeInBytes(instance.getPuppet().getContext()) + stmt.computeSizeInBytes(instance.getPuppet().getContext()) + 512 > 51200) {
                    instance.getPuppet().execute(batchStatement);
                    batchStatement = BatchStatement.newInstance(BatchType.LOGGED);
                }
                batchStatement = batchStatement.add(stmt);
            }
            if (batchStatement.size() > 0) {
                instance.getPuppet().execute(batchStatement);
            }
            long e = System.currentTimeMillis();

            logger.debug("async execution task[{}] finished, elapsed: {}ms", id, e - s);
        }));
        logger.debug("submit async execution task[{}]", id);
    }

    public static void asyncExecute(Collection<String> modifiedTables, Statement stmt) {
        int id = nextExecutionId();
        asyncExecutionFutureQueue.offer(executorService.submit(() -> {
            long s = System.currentTimeMillis();
            ZwsCassTemplateImpl.addModifiedTables(modifiedTables);
            instance.getPuppet().execute(stmt);
            logger.debug("async execution task[{}] finished, elapsed: {}ms", id, System.currentTimeMillis() - s);
        }));
        logger.debug("submit async execution task[{}]", id);
    }

    public static void asyncExecute(Collection<String> modifiedTables, Runnable runnable) {
        int id = nextExecutionId();
        asyncExecutionFutureQueue.offer(executorService.submit(() -> {
            long s = System.currentTimeMillis();
            ZwsCassTemplateImpl.addModifiedTables(modifiedTables);
            runnable.run();
            logger.debug("async execution task[{}] finished, elapsed: {}ms", id, System.currentTimeMillis() - s);
        }));
        logger.debug("submit async execution task[{}]", id);
    }

    public static void asyncExecute(String modifiedTable, String cql) {
        asyncExecute(Collections.singletonList(modifiedTable), SimpleStatement.newInstance(cql));
    }

    public static void waitAsyncExecution() {
        Future<?> future = asyncExecutionFutureQueue.poll();
        if (future != null) {
            logger.debug("enter waitAsyncExecution()");
            long s = System.currentTimeMillis();
            do {
                try {
                    future.get();
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                } catch (ExecutionException e) {
                    if (e.getCause() instanceof RuntimeException) {
                        throw (RuntimeException) e.getCause();
                    } else {
                        throw new RuntimeException(e.getCause());
                    }
                }
            } while((future = asyncExecutionFutureQueue.poll()) != null);
            logger.debug("exit waitAsyncExecution(), elapsed: {}ms", System.currentTimeMillis() - s);
        }
    }

    private static <RequestT> void throwExceptionIfNeeded(RequestT request) {
        if (needThrowException) {
            if (targetTableToThrowException == null && targetDmlTypeToThrowException == null) {
                throwException();
            }
            List<ZwsCqlInfo> infos = extractCtCqlInfos(request);
            for (ZwsCqlInfo info : infos) {
                if (targetTableToThrowException != null && targetDmlTypeToThrowException != null) {
                    if (Objects.equals(info.getTableName(), targetTableToThrowException)
                            && Objects.equals(info.getDmlType(), targetDmlTypeToThrowException)) {
                        throwException();
                    }
                    return;
                }
                if (targetTableToThrowException != null) {
                    if (Objects.equals(info.getTableName(), targetTableToThrowException)) {
                        throwException();
                    }
                    return;
                }
                if (targetDmlTypeToThrowException != null) {
                    if (Objects.equals(info.getDmlType(), targetDmlTypeToThrowException)) {
                        throwException();
                    }
                    return;
                }
            }
        }
    }

    private static void throwException() {
        if (targetExceptionToThrow != null) {
            throw targetExceptionToThrow;
        } else {
            throw EX_NO_NODE_AVAILABLE;
        }
    }

    private static <RequestT> List<ZwsCqlInfo> extractCtCqlInfos(RequestT request) {
        List<ZwsCqlInfo> infos = new ArrayList<>();
        if (request instanceof String) {
            infos.add(ZwsCqlExtractor.extract((String) request));

        } else if (request instanceof SimpleStatement) {
            infos.add(ZwsCqlExtractor.extract(((SimpleStatement) request).getQuery()));

        } else if (request instanceof BoundStatement) {
            infos.add(ZwsCqlExtractor.extract(((BoundStatement) request).getPreparedStatement().getQuery()));

        } else if (request instanceof PrepareRequest) {
            //

        } else if (request instanceof BatchStatement) {
            for (BatchableStatement<?> r : (BatchStatement) request) {
                infos.addAll(extractCtCqlInfos(r));
            }

        }
        return infos;
    }

    public static <RequestT> ZwsCqlInfo markModifiedTables(RequestT request) {
        ZwsCqlInfo info = null;
        if (request instanceof String) {
            info = ZwsCqlExtractor.extract((String) request);

        } else if (request instanceof SimpleStatement) {
            info = ZwsCqlExtractor.extract(((SimpleStatement) request).getQuery());

        } else if (request instanceof BoundStatement) {
            info = ZwsCqlExtractor.extract(((BoundStatement) request).getPreparedStatement().getQuery());

        } else if (request instanceof PrepareRequest) {
            //

        } else if (request instanceof BatchStatement) {
            for (BatchableStatement<?> r : (BatchStatement) request) {
                markModifiedTables(r);
            }

        } else {
            throw new IllegalArgumentException("unrecognized request: " + request + ", class: " + request.getClass());
        }
        logger.debug("ZwsCqlInfo: {}", info);
        if (info != null && (info.getDmlType() == ZwsCqlDmlType.Insert || info.getDmlType() == ZwsCqlDmlType.Update)) {
            ZwsCassTemplateImpl.addModifiedTable(info.getTableName());
        }
        return info;
    }

    /* don't use */
    public static void startBatching() {
        logger.debug("startBatching");
        assert instance != null;
        batching = true;
        batchStmt = BatchStatement.newInstance(BatchType.LOGGED);
    }

    /* don't use */
    public static void endBatching() {
        logger.debug("endBatching");
        assert instance != null;
        try {
            instance.getPuppet().execute(batchStmt);
        } finally {
            batching = false;
            batchStmt = null;
        }
    }

    public static void resetBatching() {
        logger.debug("resetBatching");
        batching = false;
        batchStmt = null;
    }

    /* don't use */
    private static <RequestT> void addBatch(RequestT request) {
        if (request instanceof String) {
            batchStmt = batchStmt.add(SimpleStatement.newInstance((String) request));

        } else if (request instanceof BatchableStatement) {
            batchStmt = batchStmt.add((BatchableStatement<?>) request);

        } else if (request instanceof BatchableStatement[]) {
            batchStmt = batchStmt.addAll((BatchableStatement<?>[]) request);

        } else if (request instanceof Iterable) {
            batchStmt = batchStmt.addAll((Iterable<? extends BatchableStatement<?>>) request);

        } else {
            throw new IllegalArgumentException("unrecognized request: " + request + ", class: " + request.getClass());

        }
        batchStmt = executeIfFull(batchStmt);
    }

    private static BatchStatement executeIfFull(BatchStatement batchStmt) {
        if (batchStmt.computeSizeInBytes(instance.getContext()) + 512 > 51200) {
            BatchStatement stmt = BatchStatement.newInstance(BatchType.LOGGED);
            for (BatchableStatement temp : batchStmt) {
                if (stmt.computeSizeInBytes(instance.getContext()) + temp.computeSizeInBytes(instance.getContext()) + 512 > 51200) {
                    instance.getPuppet().execute(stmt);
                    stmt = BatchStatement.newInstance(BatchType.LOGGED);
                }
                stmt = stmt.add(temp);
            }
            return stmt;
        }
        return batchStmt;
    }

    public static ZwsCqlSessionProxy create(CqlSession puppet, boolean clearEachCase) {
        instance = new ZwsCqlSessionProxy(puppet, clearEachCase);
        return instance;
    }

    private CqlSession puppet;
    private boolean clearEachCase;

    public ZwsCqlSessionProxy(CqlSession puppet, boolean clearEachCase) {
        this.puppet = puppet;
        this.clearEachCase = clearEachCase;
        ZwsCqlSessionProxyHelper.init1(puppet);
    }

    public CqlSession getPuppet() {
        return puppet;
    }

    @NonNull
    @Override
    public String getName() {
//        throwExceptionIfNeeded();
        return puppet.getName();
    }

    @NonNull
    @Override
    public Metadata getMetadata() {
//        throwExceptionIfNeeded();
        return puppet.getMetadata();
    }

    @Override
    public boolean isSchemaMetadataEnabled() {
//        throwExceptionIfNeeded();
        return puppet.isSchemaMetadataEnabled();
    }

    @NonNull
    @Override
    public CompletionStage<Metadata> setSchemaMetadataEnabled(@Nullable Boolean newValue) {
//        throwExceptionIfNeeded();
        return puppet.setSchemaMetadataEnabled(newValue);
    }

    @NonNull
    @Override
    public CompletionStage<Metadata> refreshSchemaAsync() {
//        throwExceptionIfNeeded();
        return puppet.refreshSchemaAsync();
    }

    @NonNull
    @Override
    public CompletionStage<Boolean> checkSchemaAgreementAsync() {
//        throwExceptionIfNeeded();
        return puppet.checkSchemaAgreementAsync();
    }

    @NonNull
    @Override
    public DriverContext getContext() {
//        throwExceptionIfNeeded();
        return puppet.getContext();
    }

    @NonNull
    @Override
    public Optional<CqlIdentifier> getKeyspace() {
//        throwExceptionIfNeeded();
        return puppet.getKeyspace();
    }

    @NonNull
    @Override
    public Optional<Metrics> getMetrics() {
//        throwExceptionIfNeeded();
        return puppet.getMetrics();
    }



    @Nullable
    @Override
    public <RequestT extends Request, ResultT> ResultT execute(@NonNull RequestT request,
        @NonNull GenericType<ResultT> resultType) {

        logger.debug("execute request: {}", requestInfo(request));
        throwExceptionIfNeeded(request);
        waitAsyncExecution();
        ZwsCqlInfo info = markModifiedTables(request);
        if (!batching) {
            /*
             * Shortcut for SELECT statement
             * Proxy return null in the condition of clearing data each case and no data insertion in this case
             *
             */
            if (clearEachCase
                    && info != null
                    && info.getDmlType() == ZwsCqlDmlType.Select
                    && !ZwsCassTemplateImpl.hasModifiedTable(puppet, info.getTableName())) {
                logger.warn("<y>shortcut for SELECT statement since no data in '{}'</y>", info.getTableName());
                return (ResultT) ResultSets.newInstance(Conversions.toResultSet(null, null, null, null));
            }
            return puppet.execute(request, resultType);
        } else {
            addBatch(request);
            return (ResultT) ResultSets.newInstance(Conversions.toResultSet(null, null, null, null));
        }
    }

    public <RequestT extends Request> String requestInfo(RequestT request) {
        if (request instanceof PrepareRequest) {
            return ((PrepareRequest) request).getQuery();
        }
        if (request instanceof BoundStatement) {
            return ((BoundStatement) request).getPreparedStatement().getQuery();
        }
        if (request instanceof SimpleStatement) {
            return ((SimpleStatement) request).getQuery();
        }
        if (request instanceof BatchStatement) {
            StringJoiner joiner = new StringJoiner(", ", "[", "]");
            for (BatchableStatement<?> stmt : (BatchStatement) request) {
                joiner.add(requestInfo(stmt));
            }
            return joiner.toString();
        }
        return request.toString();
    }

    @NonNull
    @Override
    public CompletionStage<Void> closeFuture() {
//        throwExceptionIfIntervened();
        return puppet.closeFuture();
    }

    @NonNull
    @Override
    public CompletionStage<Void> closeAsync() {
//        throwExceptionIfIntervened();
        return puppet.closeAsync();
    }

    @NonNull
    @Override
    public CompletionStage<Void> forceCloseAsync() {
//        throwExceptionIfIntervened();
        return puppet.forceCloseAsync();
    }
}
